import time
from locust import HttpUser, task, between
import json
from pathlib import Path
import os
import glob

'''
locust --csv=reports/locust_stats.csv --html=reports/locust_report.html -f locustfile.py --host=http://example.com

locust -f stock_center.py --no-web -c 200 -r 20 -t 1m
（–no-web 代表不需要启动UI页面
-u 代表需要并发的用户数
-r 代表每秒并发的用户数
 -t 代表需要运行的时间）
'''


def load(file, **kwargs):
    with open(file, mode="r", encoding="utf-8") as fp:
        return json.load(fp)


class QuickstartUser(HttpUser):
    wait_time = between(0.01, 0.08)

    host = "http://127.0.0.1:8000"

    def __init__(self, *args, **kwargs):
        super(QuickstartUser, self).__init__(*args, **kwargs)
        self.file = os.environ.get("test_file_param")
        self.params_request = load(self.file)

    @task
    def request_vllm(self):
        url = r"/v1/chat/completions"
        params = self.params_request
        self.client.post(url=url, json=params)


if __name__ == '__main__':
    # csv_report = "test.csv"
    # html_report = "test.html"
    # action_file = "locust_test.py"
    # c_nums = 20
    # r_nums = 10
    # action = f'locust --csv=reports/{csv_report} --html=reports/{html_report} -f {action_file} -u {c_nums} -r {r_nums} -t 1m'

    import time
    files = list(glob.glob(r"params/*.json"))
    #for f in files:
    for i in [1, 10, 20, 50, 100, 150, 200, 300]:
        for f in files:
            os.environ.setdefault("test_file_param", str(f))
            file_name = Path(f).name.split(".")[0]
            time.sleep(10)
        #for i in [1, 10, 20, 50, 100, 150, 200, 300]:
            c_nums = i * 3
            r_nums = i
            csv_report = f"{file_name}_{i}.csv"
            html_report = f"{file_name}_{i}.html"
            action_file = "locust_test.py"
            action = f'locust --csv=reports/{csv_report} --html=reports/{html_report} -f {action_file} --headless -u {c_nums} -r {r_nums} -t 2m'
            print(action)
            os.system(action)
        
